import numpy as np
import os
from PIL import Image
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
from sklearn.metrics import classification_report, accuracy_score
import matplotlib.pyplot as plt


def load_orl_faces(data_path, num_persons=40, num_images=10, img_size=(112, 92)):
    total_images = num_persons * num_images
    X = np.zeros((total_images, img_size[0] * img_size[1]))
    y = np.zeros(total_images, dtype=int)

    image_count = 0
    for person_id in range(1, num_persons + 1):
        person_dir = os.path.join(data_path, f's{person_id}')

        # 检查文件夹是否存在
        if not os.path.exists(person_dir):
            print(f"警告: 文件夹 {person_dir} 不存在，跳过")
            continue

        for img_num in range(1, num_images + 1):
            img_path = os.path.join(person_dir, f'{img_num}.pgm')

            # 检查文件是否存在
            if not os.path.exists(img_path):
                print(f"警告: 图像文件 {img_path} 不存在，跳过")
                continue

            try:
                # 读取图像并转换为numpy数组
                img = Image.open(img_path)
                img_array = np.array(img).flatten()

                X[image_count] = img_array
                y[image_count] = person_id - 1  # 标签从0开始
                image_count += 1
            except Exception as e:
                print(f"错误: 无法读取图像 {img_path}: {e}")

    # 如果实际读取的图像数量少于预期，调整数组大小
    if image_count < total_images:
        X = X[:image_count]
        y = y[:image_count]

    return X, y, img_size


class CustomPCA:
    """自定义PCA实现"""

    def __init__(self, n_components=None, whiten=True):
        self.n_components = n_components
        self.whiten = whiten
        self.components_ = None
        self.explained_variance_ratio_ = None
        self.mean_ = None
        self.explained_variance_ = None

    def fit(self, X):
        """拟合PCA模型"""
        # 计算均值
        self.mean_ = np.mean(X, axis=0)

        # 中心化数据
        X_centered = X - self.mean_

        # 计算协方差矩阵
        cov_matrix = np.cov(X_centered, rowvar=False)

        # 计算特征值和特征向量
        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)

        # 按特征值降序排序
        idx = np.argsort(eigenvalues)[::-1]
        eigenvalues = eigenvalues[idx]
        eigenvectors = eigenvectors[:, idx]

        # 确定要保留的主成分数量
        if self.n_components is not None:
            if 0 < self.n_components < 1:
                # 按方差解释比例选择成分
                total_variance = np.sum(eigenvalues)
                explained_variance_ratio = eigenvalues / total_variance
                cumulative_variance = np.cumsum(explained_variance_ratio)
                self.n_components_ = np.argmax(cumulative_variance >= self.n_components) + 1
            else:
                self.n_components_ = min(self.n_components, eigenvectors.shape[1])
        else:
            self.n_components_ = eigenvectors.shape[1]

        # 保存主成分
        self.components_ = eigenvectors[:, :self.n_components_].T
        self.explained_variance_ = eigenvalues[:self.n_components_]
        self.explained_variance_ratio_ = self.explained_variance_ / np.sum(eigenvalues)

        return self

    def transform(self, X):
        """将数据转换到主成分空间"""
        X_centered = X - self.mean_
        X_transformed = np.dot(X_centered, self.components_.T)

        # 白化处理（可选）
        if self.whiten:
            X_transformed /= np.sqrt(self.explained_variance_)

        return X_transformed

    def fit_transform(self, X):
        """拟合并转换数据"""
        return self.fit(X).transform(X)


def main():
    # 1. 加载数据
    data_path = "orl_faces"  # 修改为您的实际路径
    print("正在加载ORL人脸数据集...")

    try:
        X, y, image_shape = load_orl_faces(data_path)
        print(f"数据集信息:")
        print(f"样本数: {X.shape[0]}, 特征数: {X.shape[1]}")
        print(f"类别数: {len(np.unique(y))}")
        print(f"图像尺寸: {image_shape}")
    except Exception as e:
        print(f"数据加载错误: {e}")
        return

    # 2. 数据预处理 - 中心化
    mean_face = np.mean(X, axis=0)
    X_centered = X - mean_face

    # 3. 划分训练集和测试集 (70%训练, 30%测试)
    X_train, X_test, y_train, y_test = train_test_split(
        X_centered, y, test_size=0.3, random_state=42, stratify=y
    )

    print(f"\n训练集大小: {X_train.shape}")
    print(f"测试集大小: {X_test.shape}")

    # 4. 使用自定义PCA降维
    print("正在进行PCA降维...")
    try:
        # 使用自定义PCA
        pca = CustomPCA(n_components=0.95, whiten=True)
        X_train_pca = pca.fit_transform(X_train)
        X_test_pca = pca.transform(X_test)

        print(f"降维后训练集形状: {X_train_pca.shape}")
        print(f"实际保留的主成分数: {pca.n_components_}")
        print(f"累计方差解释率: {np.sum(pca.explained_variance_ratio_):.3f}")
    except Exception as e:
        print(f"PCA降维错误: {e}")
        # 备用方案：使用固定数量的主成分
        n_components = min(150, X_train.shape[0], X_train.shape[1])
        print(f"尝试使用固定主成分数: {n_components}")
        pca = CustomPCA(n_components=n_components, whiten=True)
        X_train_pca = pca.fit_transform(X_train)
        X_test_pca = pca.transform(X_test)

    # 5. 训练SVM分类器
    print("训练SVM分类器...")
    svm_classifier = SVC(kernel='rbf', class_weight='balanced', random_state=42)
    svm_classifier.fit(X_train_pca, y_train)

    # 6. 预测与评估
    y_pred = svm_classifier.predict(X_test_pca)
    accuracy = accuracy_score(y_test, y_pred)

    print(f"\n模型性能评估:")
    print(f"识别准确率: {accuracy:.3f}")
    print("\n详细分类报告:")
    print(classification_report(y_test, y_pred, zero_division=0))

    # 7. 可视化特征脸（Eigenfaces）
    try:
        eigenfaces = pca.components_
        fig, axes = plt.subplots(4, 4, figsize=(12, 12))
        fig.suptitle('前16个特征脸 (Eigenfaces)', fontsize=16)

        for i, ax in enumerate(axes.flat):
            if i < min(16, pca.n_components_):
                eigenface = eigenfaces[i].reshape(image_shape)
                ax.imshow(eigenface, cmap='gray')
                ax.set_title(f'主成分 {i + 1}')
                ax.axis('off')

        plt.tight_layout()
        plt.show()

        # 8. 可视化均值脸
        plt.figure(figsize=(8, 6))
        mean_face_img = mean_face.reshape(image_shape)
        plt.imshow(mean_face_img, cmap='gray')
        plt.title('均值脸 (Mean Face)')
        plt.axis('off')
        plt.show()

        # 9. 可视化部分测试结果
        X_test_original = X_test + mean_face  # 还原中心化前的数据
        n_samples = min(12, len(X_test))
        indices = np.random.choice(len(X_test), n_samples, replace=False)

        fig, axes = plt.subplots(3, 4, figsize=(15, 12))
        fig.suptitle('测试集识别结果示例', fontsize=16)

        for i, ax in enumerate(axes.flat):
            if i < n_samples:
                idx = indices[i]
                img = X_test_original[idx].reshape(image_shape)
                ax.imshow(img, cmap='gray')
                color = 'green' if y_test[idx] == y_pred[idx] else 'red'
                ax.set_title(f'真实: {y_test[idx] + 1}\n预测: {y_pred[idx] + 1}',
                             color=color, fontweight='bold')
                ax.axis('off')

        plt.tight_layout()
        plt.show()

    except Exception as e:
        print(f"可视化错误: {e}")


if __name__ == "__main__":
    main()